import os

from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess

from ai_core.asr.base import Asr
from ai_core.audio_format.opus import OpusEncode
from ai_core.utils.util import Util


class FunAsr(Asr):
    def __init__(self,config,model_dir=None):
        self.output_dir = Util.get_project_dir() + config.get("output_dir")
        print(self.output_dir)
        self.model_dir = model_dir or config.get("model_dir")
        # 当前文件做测试时，需要使用绝对路径，配置文件中的路径会导致找不到文件
        # self.model_dir = "D:\\code\\ai_test\\model\\fun_asr"
        print(self.model_dir)
        self.model = AutoModel(
            model = self.model_dir,
            vad_kwargs = {"max_single_segment_time": 30000},
            # device = "cuda:0", # 默认 device = "cpu"
            disable_update = True,
            hub="hf",
        )
        # 测试
        # self.raw_outputs = {} # 用于存储模型原始输出
        # self.hooks = [] # 用于存储模型 hook

    # 语音识别,将音频文件转化为文字
    def audio_file_to_text(self, file_path):
        res_text = self.model.generate(
            input = file_path,
            language = "auto",
            use_itn = True, # 是否使用 ITN, 归一化（加标点，将三点五转换为3.5）
            batch_size = 60, # 一次预测的音频数量
            merge_va = True, # 是否合并音节
            merge_length_s = 100, # 合并音节的最小长度
        )
        # return res[0]["text"]
        return rich_transcription_postprocess(res_text[0]["text"])

    def opus_data_to_text(self, opus_data, format_str="wav"):
        # opus = OpusEncode()
        # opus.opus_to_wav_file(file_path, opus_data)
        file_path = Util.get_random_file_path(self.output_dir, format_str)
        print(f"[FunAsr] opus_data_to_text:  {file_path}")
        super().opus_data_to_wav(opus_data, file_path)
        # 判断输出文件是否存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件 {file_path} 不存在")
        return self.audio_file_to_text(file_path)
        # return "test"

    # 测试
    # def _register_hook(self):
    #     """注册钩子到输出层（示例： Paraformer 的 CTC头部）"""
    #     # 假设输出曾是 mode了.ctc.ctc_lo
    #     output_layer = self.model.model.ctc.ctc_lo
    #
    #     # 定义钩子函数
    #     def hook_function(module, inputs, outputs):
    #         # output 是logits（为归一化的概率）
    #         self.raw_outputs["ctc_logits"] = outputs.clone()
    #
    #     # 注册钩子
    #     self.hooks.append(output_layer.register_forward_hook(hook_function))
    #
    # def _remove_hook(self):
    #     for hook in self.hooks:
    #         hook.remove()
    #     self.hooks.clear()
    #     self.raw_outputs.clear()



if __name__ == "__main__":
    # fun_asr = FunAsr()
    # fun_asr._register_hook()
    # res = fun_asr.audio_file_to_text("test.wav")
    #
    # # 打印模型原始输出
    # if fun_asr.raw_outputs:
    #     logits = fun_asr.raw_outputs["ctc_logits"]
    #     print("模型传给分词器的原始 Logits: 形状", logits.shape)
    #
    #     # 获取 Tokens IDs (假设使用 argmax 解码)
    #     tokens_ids = logits.argmax(dim=-1).squeeze().tolist()
    #     print("Tokens IDs: ", tokens_ids)
    #
    # print(res)

    asr_config = Util.get_config()
    # 直接通过参数传递（需类支持）
    fun_asr = FunAsr(
        asr_config.get("ASR").get("Fun_ASR"),  # 原始配置
        model_dir="D:\\code\\ai_test\\model\\fun_asr"  # 显式指定路径
    )

    opusEncode = OpusEncode()
    opus_data = opusEncode.load_opus_custom_data(Util.get_project_dir() + "test.opus")

    resource = fun_asr.opus_data_to_text(opus_data)

    print(resource)